package gateway

import (
	"azfirewall-gateway/service"
	"net/http"
	"strings"
)

func rewriteResponse(resp *http.Response) (err error) {
	r := resp.Request
	var site = service.GetSite(r.Host)
	locationURL, err := resp.Location()
	if locationURL != nil {
		port := locationURL.Port()
		if (port != "80") && (port != "443") {
			host := locationURL.Hostname()
			if site != nil {
				newLocation := strings.Replace(locationURL.String(), host+":"+port, host, -1)
				userScheme := "http"
				if resp.Request.TLS != nil {
					userScheme = "https"
				}
				newLocation = strings.Replace(newLocation, locationURL.Scheme, userScheme, 1)
				resp.Header.Set("Location", newLocation)
			}
		}
	}
	resp.Header.Set("X-WAF-Designed-By", "Tony_An || Starhome_Cloud")
	resp.Header.Set("X-Powered-By", "CDNPlus")
	//TODO CC防护部分内容

	// HSTS
	if resp.Request.TLS != nil && site.HSTSEnabled {
		resp.Header.Set("Strict-Transport-Security", "max-age=31536000; includeSubDomains")
	}
	return nil
}
